
import sys,matplotlib
from matplotlib import colors, ticker, cm
sys.path.append('../../utils/')
matplotlib.use('Agg')
from numpy import *
from pylab import *
import getopt
import os, re
import pdb
import math as math
import linecache
from matplotlib.colors import LogNorm
from scipy import interpolate
import pickle as pickle


matplotlib.rc('font', family='serif')
matplotlib.rc('xtick', labelsize = 15)
matplotlib.rc('ytick', labelsize = 15)
matplotlib.rc('axes', linewidth=2)
fontsize = 20
fontname = 'serif'
linewidth = 2
annosize=15
annoweight=300
labelpad=5

#fname = '/net/wukong.gps.caltech.edu/ssd1/incoming/willacy/pluto_4/titan/1d/pluto_3_gas_ss_flare/gas2022.out-100'
#fname = 'gas2022.out-100_flare'
fname = 'kinmars-cole.out'

infile = open(fname)
nz = 13

dat = {}
tctr = 0
#ntime = 10930
ntime =74 # 74

temp = infile.readline()
while not 'TSTEP' in temp:
    print(temp)
    temp = infile.readline()
tcheck = temp
print(tcheck)

while 'SEC' in tcheck: # make sure it's the TSTEP 
    tctr += 1
    print(tctr)
    dat[tctr] = {}
    if tctr == 1:
        dat[tctr]['time'] = float(tcheck.split()[3])
    else:
        dat[tctr]['time'] = dat[tctr-1]['time'] + float(tcheck.split()[3])

    while 'MIXING' not in temp:
        temp = infile.readline()
    temp = infile.readline() # blank
    header = infile.readline().split() # species
    while len(header) >1:
        for ix in range(1,len(header)):
            dat[tctr][header[ix]] = np.zeros(nz)
        alt = np.zeros(nz)
        for iz in range(0,nz):
            temp = infile.readline().split()
            for ix in range(1,len(header)):
                dat[tctr][header[ix]][iz] = float(temp[ix+1])
                alt[iz] = float(temp[1])
        dat[tctr]['alt'] = alt
        temp = infile.readline() # blank                                                                                                                                                                   
        header = infile.readline().split()
    if tctr < ntime:
        while (not 'TSTEP' in temp):
            temp = infile.readline()
        tcheck = temp
    else:
        break


########## ########## ########## ########## ########## ########## ########## ########## 

t, CH4, C2H2, C2H4, C2H6, HCN = np.zeros((len(list(dat.keys())))), np.zeros((len(list(dat.keys())),nz)),\
                                np.zeros((len(list(dat.keys())),nz)),np.zeros((len(list(dat.keys())),nz)),\
                                np.zeros((len(list(dat.keys())),nz)),np.zeros((len(list(dat.keys())),nz))
C4H2, C6H2, C8H2, hcnh, n = np.zeros((len(list(dat.keys())),nz)), np.zeros((len(list(dat.keys())),nz)),\
                            np.zeros((len(list(dat.keys())),nz)),np.zeros((len(list(dat.keys())),nz)), np.zeros((len(list(dat.keys())),nz))

print(np.shape(CH4))
ctr = 0
for i in list(dat.keys()):
    t[ctr] = float(dat[i]['time'])*ctr
    CH4[ctr,:] = np.array(dat[i]['H2'])
#    C2H2[ctr,:] = np.array(dat[i]['C2H2'])
#    C2H4[ctr,:] = np.array(dat[i]['C2H4'])
#    C2H6[ctr,:] = np.array(dat[i]['C2H6'])
#    HCN[ctr,:] = np.array(dat[i]['HCN'])
#    C4H2[ctr,:] = np.array(dat[i]['C4H2'])
#    C6H2[ctr,:] = np.array(dat[i]['C6H2'])
#    C8H2[ctr,:] = np.array(dat[i]['C8H2'])
#    hcnh[ctr,:] = np.array(dat[i]['HCNH+'])
#    n[ctr,:] = np.array(dat[i]['N+'])

    ctr += 1

#################### ########## ########## ########## ########## ########## ########## 

fig,ax=plt.subplots(figsize=(8,3))
print(t)
ax.semilogx(t[:-1],CH4[:-1,0],'blue')
ax.set_xlim([1e5,1e18])
plt.savefig('timeh2.png',dpi=300)
exit()

# severely hardcoded, apologies to future self

fig,ax=plt.subplots(nrows=6,ncols=1,figsize=(9,24))
#30, 100, 200, 500, 1000 km = lvl 8, 11, 15, 24, 36

t[-1]=t[-2]

alts = [10,18,24,32,34,39]
labels = ['100 km','300 km','500 km','800 km','1000 km','1300 km']


# add C4H2, C6H2, C8H2
# 0 km 
ax[0].semilogx(t,CH4[:,alts[0]]/CH4[0,alts[0]],'grey',label='CH4',linewidth=linewidth)
ax[0].semilogx(t,C2H2[:,alts[0]]/C2H2[0,alts[0]],'red',label='C2H2',linestyle='-',linewidth=linewidth)
ax[0].semilogx(t,C2H4[:,alts[0]]/C2H4[0,alts[0]],'purple',label='C2H4',linestyle='-',linewidth=linewidth)
ax[0].semilogx(t,C2H6[:,alts[0]]/C2H6[0,alts[0]],'blue',label='C2H6',linestyle='-',linewidth=linewidth)
ax[0].semilogx(t,HCN[:,alts[0]]/HCN[0,alts[0]],'green',label='HCN',linestyle='-',linewidth=linewidth)
ax[0].semilogx(t,C4H2[:,alts[0]]/C4H2[0,alts[0]],'orange',label='C4H2',linestyle='-',linewidth=linewidth)
ax[0].semilogx(t,C6H2[:,alts[0]]/C6H2[0,alts[0]],'cyan',label='C6H2',linestyle='-',linewidth=linewidth)
ax[0].semilogx(t,C8H2[:,alts[0]]/C8H2[0,alts[0]],'magenta',label='C8H2',linestyle='-',linewidth=linewidth)
#ax[0].semilogx(t,hcnh[:,alts[0]]/hcnh[0,alts[0]],'lime',label='HCNH+',linestyle='-',linewidth=linewidth)
#ax[0].semilogx(t,n[:,alts[0]]/n[0,alts[0]],'pink',label='N+',linestyle='-',linewidth=linewidth)
ax[0].set_title(labels[0],fontsize=fontsize,fontname=fontname)


# 30 km , lvl #8, or indx #7
ax[1].semilogx(t,CH4[:,alts[1]]/CH4[0,alts[1]],'grey',label='CH4',linewidth=linewidth)
ax[1].semilogx(t,C2H2[:,alts[1]]/C2H2[0,alts[1]],'red',label='C2H2',linestyle='-',linewidth=linewidth)
ax[1].semilogx(t,C2H4[:,alts[1]]/C2H4[0,alts[1]],'purple',label='C2H4',linestyle='-',linewidth=linewidth)
ax[1].semilogx(t,C2H6[:,alts[1]]/C2H6[0,alts[1]],'blue',label='C2H6',linestyle='-',linewidth=linewidth)
ax[1].semilogx(t,HCN[:,alts[1]]/HCN[0,alts[1]],'green',label='HCN',linestyle='-',linewidth=linewidth)
ax[1].semilogx(t,C4H2[:,alts[1]]/C4H2[0,alts[1]],'orange',label='C4H2',linestyle='-',linewidth=linewidth)
ax[1].semilogx(t,C6H2[:,alts[1]]/C6H2[0,alts[1]],'cyan',label='C6H2',linestyle='-',linewidth=linewidth)
ax[1].semilogx(t,C8H2[:,alts[1]]/C8H2[0,alts[1]],'magenta',label='C8H2',linestyle='-',linewidth=linewidth)
#ax[1].semilogx(t,hcnh[:,alts[1]]/hcnh[0,alts[1]],'lime',label='HCNH+',linestyle='-',linewidth=linewidth)
#ax[1].semilogx(t,n[:,alts[1]]/n[0,alts[1]],'pink',label='N+',linestyle='-',linewidth=linewidth)
ax[1].set_title(labels[1],fontsize=fontsize,fontname=fontname)

# 100 km, lvl #11 or indx #10
ax[2].semilogx(t,CH4[:,alts[2]]/CH4[0,alts[2]],'grey',label='CH4',linewidth=linewidth)
ax[2].semilogx(t,C2H2[:,alts[2]]/C2H2[0,alts[2]],'red',label='C2H2',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,C2H4[:,alts[2]]/C2H4[0,alts[2]],'purple',label='C2H4',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,C2H6[:,alts[2]]/C2H6[0,alts[2]],'blue',label='C2H6',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,HCN[:,alts[2]]/HCN[0,alts[2]],'green',label='HCN',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,C4H2[:,alts[2]]/C4H2[0,alts[2]],'orange',label='C4H2',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,C6H2[:,alts[2]]/C6H2[0,alts[2]],'cyan',label='C6H2',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,C8H2[:,alts[2]]/C8H2[0,alts[2]],'magenta',label='C8H2',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,hcnh[:,alts[2]]/hcnh[0,alts[2]],'lime',label='HCNH+',linestyle='-',linewidth=linewidth)
ax[2].semilogx(t,n[:,alts[2]]/n[0,alts[2]],'pink',label='N+',linestyle='-',linewidth=linewidth)
ax[2].set_title(labels[2],fontsize=fontsize,fontname=fontname)

# 200 km , lvl #15 or indx 14
ax[3].semilogx(t,CH4[:,alts[3]]/CH4[0,alts[3]],'grey',label='CH4',linewidth=linewidth)
ax[3].semilogx(t,C2H2[:,alts[3]]/C2H2[0,alts[3]],'red',label='C2H2',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,C2H4[:,alts[3]]/C2H4[0,alts[3]],'purple',label='C2H4',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,C2H6[:,alts[3]]/C2H6[0,alts[3]],'blue',label='C2H6',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,HCN[:,alts[3]]/HCN[0,alts[3]],'green',label='HCN',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,C4H2[:,alts[3]]/C4H2[0,alts[3]],'orange',label='C4H2',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,C6H2[:,alts[3]]/C6H2[0,alts[3]],'cyan',label='C6H2',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,C8H2[:,alts[3]]/C8H2[0,alts[3]],'magenta',label='C8H2',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,hcnh[:,alts[3]]/hcnh[0,alts[3]],'lime',label='HCNH+',linestyle='-',linewidth=linewidth)
ax[3].semilogx(t,n[:,alts[3]]/n[0,alts[3]],'pink',label='N+',linestyle='-',linewidth=linewidth)
ax[3].set_title(labels[3],fontsize=fontsize,fontname=fontname)

# 500 km, lvl #24 or indx 23
ax[4].semilogx(t,CH4[:,alts[4]]/CH4[0,alts[4]],'grey',label='CH4',linewidth=linewidth)
ax[4].semilogx(t,C2H2[:,alts[4]]/C2H2[0,alts[4]],'red',label='C2H2',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,C2H4[:,alts[4]]/C2H4[0,alts[4]],'purple',label='C2H4',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,C2H6[:,alts[4]]/C2H6[0,alts[4]],'blue',label='C2H6',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,HCN[:,alts[4]]/HCN[0,alts[4]],'green',label='HCN',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,C4H2[:,alts[4]]/C4H2[0,alts[4]],'orange',label='C4H2',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,C6H2[:,alts[4]]/C6H2[0,alts[4]],'cyan',label='C6H2',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,C8H2[:,alts[4]]/C8H2[0,alts[4]],'magenta',label='C8H2',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,hcnh[:,alts[4]]/hcnh[0,alts[4]],'lime',label='HCNH+',linestyle='-',linewidth=linewidth)
ax[4].semilogx(t,n[:,alts[4]]/n[0,alts[4]],'pink',label='N+',linestyle='-',linewidth=linewidth)
ax[4].set_title(labels[4],fontsize=fontsize,fontname=fontname)

# 1000 km, lvl alts[2] or indx alts[5]
ax[5].semilogx(t,CH4[:,alts[5]]/CH4[0,alts[5]],'grey',label='CH4',linewidth=linewidth)
ax[5].semilogx(t,C2H2[:,alts[5]]/C2H2[0,alts[5]],'red',label='C2H2',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,C2H4[:,alts[5]]/C2H4[0,alts[5]],'purple',label='C2H4',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,C2H6[:,alts[5]]/C2H6[0,alts[5]],'blue',label='C2H6',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,HCN[:,alts[5]]/HCN[0,alts[5]],'green',label='HCN',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,C4H2[:,alts[5]]/C4H2[0,alts[5]],'orange',label='C4H2',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,C6H2[:,alts[5]]/C6H2[0,alts[5]],'cyan',label='C6H2',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,C8H2[:,alts[5]]/C8H2[0,alts[5]],'magenta',label='C8H2',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,hcnh[:,alts[5]]/hcnh[0,alts[5]],'lime',label='HCNH+',linestyle='-',linewidth=linewidth)
ax[5].semilogx(t,n[:,alts[5]]/n[0,alts[5]],'pink',label='N+',linestyle='-',linewidth=linewidth)
ax[5].set_title(labels[5],fontsize=fontsize,fontname=fontname)

ax[2].set_ylabel('mr(t)/mr(t=0)',fontsize=fontsize,fontname=fontname)
ax[5].set_xlabel('time in seconds',fontsize=fontsize,fontname=fontname)
ax[1].legend()

plt.subplots_adjust(hspace=0.35)

for ix in range(0,6):
    ax[ix].set_xlim([1e5,1e12])
    ax[ix].set_ylim([0.7,2.0])
    ax[ix].tick_params('both', length=10, width=2, which='major')
    ax[ix].tick_params('both', length=5, width=1, which='minor')
    ax[ix].tick_params(axis='x', pad=8)
    ax[ix].tick_params(axis='y', pad=8)

plt.savefig('tseries_1.25_log.png')
#plt.savefig('tseries_srmax.png')
#plt.savefig('tseries_2003_short_log.png')
